import sys 
sys.path.append('/home/aistudio/external-libraries')

import numpy as np
import meshio
from scipy.spatial import Delaunay
from meshpy.tet import MeshInfo, build, Options
from restore_model_tools import restore_model, plot_room_flow_field, standardize_NS_Equations_3D, data
from mesh_refinement_tools import plot_room_add_points_3d, get_top_k_indices, tetrahedron_centroid, tetrahedron_midpoint, remove_duplicate_points

# 对哪个level的网格进行加密？
level = 2

# 加载模型
level2_room_model = restore_model(model_path = "../2_PINN_training/Model/model-70000.pdparams")
# plot_room_flow_field(level1_room_model)

# Get the node coordinates
mesh = meshio.read('../0_Mesh_files/level2_refined_room.vtk', file_format='vtk')
points = mesh.points

# Delaunay四面体化
# points 输入散点；furthest_site 为True时，计算最远点；incremental 为True时，允许增量添加点；qhull_options 为qhull参数，具体可参考qhull
tri = Delaunay(points, furthest_site=False, incremental=False, qhull_options=None)
# 加密前：获取三角形网格的顶点和三角形索引
vertices = tri.points
tetrahedrons = tri.simplices

# 计算每个四面体的质心
centroid_points = tetrahedron_centroid(tetrahedrons, vertices)
X = centroid_points
print("原网格四面体个数:", X.shape)

# 找到残差较大的质心点
standardize_X = data.transform_inputs(X)
[f1, f2, f3, f4] = level2_room_model.predict(standardize_X, operator=standardize_NS_Equations_3D)
err_eq = np.absolute(f1) + np.absolute(f2) + np.absolute(f3) + np.absolute(f4)
x_id = get_top_k_indices(err_eq, k=1000)
add_points = X[x_id]
plot_room_add_points_3d(add_points)

# 计算需要被细化的单元的Edge中点
tets_id = x_id
refine_tets = tetrahedrons[tets_id]
refine_midpoints = tetrahedron_midpoint(refine_tets, vertices)
refine_midpoints = remove_duplicate_points(refine_midpoints)
print("待加密四面体网格个数:", refine_tets.shape)
print("待加密四面体边中点个数:", refine_midpoints.shape)
plot_room_add_points_3d(refine_midpoints)

# 对加密后的原有网格节点进行四面体化并输出VTK文件
# points 输入散点；furthest_site 为True时，计算最远点；incremental 为True时，允许增量添加点；qhull_options 为qhull参数，具体可参考qhull
refined_tet = Delaunay(np.vstack((vertices, refine_midpoints)), furthest_site=False, incremental=False, qhull_options=None)
refined_mesh_info = MeshInfo()
refined_mesh_info.set_points(refined_tet.points)
refined_mesh_info.set_facets(refined_tet.convex_hull)
# mesh_info.set_elements(refined_tet.simplices)
refined_mesh = build(refined_mesh_info, options=Options("p"))
refined_mesh.write_vtk("level" + str(level + 1) + "_refined_room.vtk")


